#lang scheme

(provide rope-fold
         rope-fold-right
         rope->list
         rope->vector
         
         brope-length
         brope-for-each
         brope-length
         write-brope
         brope->bytes!
         brope->bytes
         brope-null?
         make-bjoin
         )

(require "test.ss"
         "bytes.ss"
         mzlib/etc
         srfi/1
         srfi/26
         srfi/43)


; TODO: include domain
; TODO: include byteslice
; TODO: include file content
; TODO: slice of brope
; TODO: mutable/immutable brope
; TODO: rope-balance 
; TODO: rope-flatten 
; TODO: rope-numsegments
; TODO: rope-depth
; TODO: rope -> anytree, brope -> banytree


(define-struct bjoin (lst ; list/vector of brope
                      inter))

(define (rope-fold proc init br)
  (cond
    [(vector? br)  (vector-fold (lambda(_ init b) (rope-fold proc init b)) init br)]
    [(null? br)    init]
    [(pair? br)    (rope-fold proc (rope-fold proc init (car br)) (cdr br))]
    [(bjoin? br)   (let ([lst (bjoin-lst br)]
                         [inter (bjoin-inter br)])
                     (cond 
                       [(vector? lst) 
                        (vector-fold 
                         (lambda(i init b)
                           (rope-fold proc 
                                      (if (zero? i) 
                                          init
                                          (rope-fold proc init inter))
                                      b))
                         init lst)]
                       [(null? lst) init]
                       [(pair? lst) 
                        (fold (lambda(b init)
                                (rope-fold proc (rope-fold proc init inter) b))
                              (rope-fold proc init (car lst))
                              (cdr lst))]
                       ))]
    [else   (proc br init)]))


(define (rope-fold-right proc init br)
  (cond
    [(vector? br)  (vector-fold-right (lambda(_ init b) (rope-fold-right proc init b ))
                                      init 
                                      br)]
    [(null? br)    init]
    [(pair? br)    (rope-fold-right proc (rope-fold-right proc init (cdr br)) (car br))]
    [(bjoin? br)   (let ([lst (bjoin-lst br)]
                         [inter (bjoin-inter br)])
                     (cond 
                       [(vector? lst) 
                        (vector-fold-right
                         (lambda(i init b)
                           (if (zero? i)
                               (rope-fold-right proc init b)
                               (rope-fold-right proc 
                                                (rope-fold-right proc init b)
                                                inter)))
                         init lst)]
                       [(null? lst) init]
                       [(pair? lst)
                        (rope-fold-right proc 
                                         (fold-right (lambda(b init)
                                                       (rope-fold-right proc (rope-fold-right proc init b) inter))
                                                     init
                                                     (cdr lst))
                                         (car lst))]
                       ))]
    [else   (proc br init)]))

(define (rope-length br)
  (rope-fold (lambda (b sum) (+ sum 1)) 0 br))

(define (rope-for-each proc br)
  (brope-fold (lambda (b i) (proc i b) (add1 i)) 0 br)
  (void))

(define (rope->list br [proc identity])
;  (rope-fold-right cons '() br))
  (rope-fold-right (lambda(b init)(cons (proc b) init)) '() br))
  
(define (rope->vector br [proc identity])
  (let ([v (make-vector (rope-length br))])
    (rope-for-each (lambda(i b)(vector-set! v i (proc b))) br)
    v))

(define (brope-fold proc init br)
  (rope-fold (lambda(b init)(proc (if(byte? b)(bytes b)b) init)) init br))

(define (brope-fold-right proc init br)
  (rope-fold-right (lambda(b init)(proc (if(byte? b)(bytes b)b) init)) init br))

(define (brope-for-each proc br)
  (brope-fold (lambda (b _) (proc b)) #f br)
  (void))

(define (brope-length-1 br)
  (brope-fold (lambda (b sum) (+ sum (bytes-length b))) 0 br))

; slow, TODO: try call-with-output-bytes or first (make-bytes brope-length), then fill
(define (brope->bytes-1 . br)
  (apply bytes-append (brope-fold-right cons '() br)))

(define (brope->bytes-2 . br)
  (call-with-output-bytes
   (cut write-brope br <>)))

(define (brope->bytes! bstr src-start . br)
    (brope-fold (lambda (b i)
                 (bytes-copy! bstr i b)
                 (+ i (bytes-length b)))
               src-start br)
    bstr)

(define (brope->bytes . br)
  (let* ([len (brope-length br)]
         [bstr (make-bytes len)])
    (brope->bytes! bstr 0 br)))

(define (write-brope br [port (current-output-port)])
  (brope-for-each (cut write-bytes <> port) br))

(define (brope-null? br)
  (or
;   [not br] ; #f
   [and (bytes? br)   (bytes-null? br)]
   [and (vector? br)  (vector-every brope-null? br)]
   [and (null? br)    #t]
   [and (pair? br)    (car br) (cdr br)]
;   [and (list? br)    (every brope-null? br)] ; TODO: pair
   [and (bjoin? br)   
        (brope-null? (bjoin-lst br))
        (brope-null? (bjoin-inter br))]))


; may be faster?
(define (brope-length br)
  (cond
    [(bytes? br) (bytes-length br)]
    [(byte? br) 1]
    [(vector? br) (vector-fold (lambda(_ sum b) (+ sum (brope-length b))) 0 br)]
    [(null? br)   0]
    [(pair? br)   (+ (brope-length (car br)) (brope-length (cdr br)))]
    [(bjoin? br) (let ([lst (bjoin-lst br)]
                       [inter (bjoin-inter br)])
                   (cond 
                     [(vector? lst) (+ (brope-length lst) (* (sub1 (vector-length lst)) (brope-length inter)))]
                     [(list? lst) (+ (brope-length lst) (* (sub1 (length lst)) (brope-length inter)))]))]
    [else (error "brope-length:" br)]
    ))

#|
(define (write-brope br [port (current-output-port)])
  (cond
    [(bytes? br) (write-bytes br port)]
    [(byte? br) (write-byte br port)]
    [(vector? br) (vector-for-each (lambda(_ b) (write-brope b port)) br)]
    [(list? br) (for-each (lambda(b) (write-brope b port)) br)]
    [(bjoin? br) (let ([lst (bjoin-lst br)]
                       [inter (bjoin-inter br)])
                   (cond 
                     [(vector? lst) 
                      (vector-for-each (lambda(i b) 
                                         (unless (zero? i) (write-brope inter port)) 
                                         (write-brope b port)) 
                                       lst)]
                     [(null? lst) ]
                     [(list? lst) 
                      (write-brope (car lst) port)
                      (for-each (lambda(b)
                                  (write-brope inter port)
                                  (write-brope b port))
                                (cdr lst))]))]
    [else (error "write-brope:" br)])
  (void))
|#
; +TEST
(define test-ropes (list #"Hello, world"
                    (vector #"Hello" #", " #"world")
                    (vector (vector #"H" #"e" #"l" #"l" #"o" )
                            #", " 
                            #"world")
                    (vector  (vector 72 101 108 108 111)
                             #", " 
                             #"world")
                    (list (list 72 101 108 108 111)
                          #", " 
                          '()
                          #"world")
                    (make-bjoin '(#"first" #"second" #"third") #", ")
                    (make-bjoin #(#"first" #"second" #"third") #", ")
                    ))

(equal! (map  (cut brope-fold cons '() <>) test-ropes)
        '((#"Hello, world")
          (#"world" #", " #"Hello")
          (#"world" #", " #"o" #"l" #"l" #"e" #"H")
          (#"world" #", " #"o" #"l" #"l" #"e" #"H")
          (#"world" #", " #"o" #"l" #"l" #"e" #"H")
          (#"third" #", " #"second" #", " #"first")
          (#"third" #", " #"second" #", " #"first")))

(equal! (map  (cut brope-fold-right cons '() <>) test-ropes)
        '((#"Hello, world")
          (#"Hello" #", " #"world")
          (#"H" #"e" #"l" #"l" #"o" #", " #"world")
          (#"H" #"e" #"l" #"l" #"o" #", " #"world")
          (#"H" #"e" #"l" #"l" #"o" #", " #"world")
          (#"first" #", " #"second" #", " #"third")
          (#"first" #", " #"second" #", " #"third")))

(equal! (map brope->bytes test-ropes) 
        '(#"Hello, world"
          #"Hello, world"
          #"Hello, world"
          #"Hello, world"
          #"Hello, world"
          #"first, second, third"
          #"first, second, third"))

(equal! (map brope-length test-ropes) 
        (map brope-length-1 test-ropes)
        (map (compose bytes-length brope->bytes) test-ropes)
        '(12 12 12 12 12 20 20))


(equal! (call-with-output-bytes
         (cut write-brope (make-bjoin test-ropes 10) <>))
        (brope->bytes-1 (make-bjoin test-ropes 10))
        (brope->bytes-2 (make-bjoin test-ropes 10))
        (brope->bytes (make-bjoin test-ropes 10))
        #"Hello, world\nHello, world\nHello, world\nHello, world\nHello, world\nfirst, second, third\nfirst, second, third")
